#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author: bao
# @Date:   2017-01-24 10:11:47
# @Last Modified by:   bao
# @Last Modified time: 2017-05-15 20:07:17

import math
import random
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime

'''
Flow generator
Static:
Server-Stride
ToR-Stride
Random
Cluster

Dynamic:
Data-mining
Web-Search
'''

def print_matrix(matrix, title='Traffic Demand Matrix', xlabel='Dst Host', ylabel='Src Host', file_name='matrix'):
    '''
    Print the traffic demand matrix in heatmap format
    '''
    plt.imshow(matrix, cmap='hot_r', origin='lower', interpolation='nearest',
        extent=(0.5,np.shape(matrix)[0]+0.5,0.5,np.shape(matrix)[1]+0.5))
    #plt.title(title)
    plt.xlabel(xlabel, fontsize=14)
    plt.ylabel(ylabel, fontsize=14)
    plt.colorbar()
    plt.savefig('%s.pdf' %file_name, bbox_inches='tight')
    plt.show()
    
def makeFlows(flows_src, flows_dst, demand_type, switch_nums, host_num, host_nums, cluster_size, threshold=0.1):

    matrix = np.zeros((host_nums,host_nums), dtype=np.float16)
    if demand_type == "random":
        #print ("random")
        for src in xrange(0,host_nums):
            dst = random.randint(0,host_nums-1)
            matrix[src][dst] = random.random()
            if matrix[src][dst]>threshold:
                if src not in flows_src:
                    flows_src[src]={}
                if dst not in flows_dst:
                    flows_dst[dst]={}
                flows_src[src][dst] = {'converged': False, 'demand': matrix[src][dst], 'recLimited': False}
                flows_dst[dst][src] = {'converged': False, 'demand': matrix[src][dst], 'recLimited': False}
    elif demand_type == "cluster":
        #print ("cluster")
        switches = list(xrange(0, switch_nums))
        #print switches
        
        while switches:
            if len(switches) > cluster_size and len(switches) < 2*cluster_size:
                cluster_i = switches
            else:
                cluster_i = random.sample(switches, cluster_size)
            #print "cluster_i", cluster_i
            switches = list(set(switches)-set(cluster_i))
            for src_sw in cluster_i:
                for x in xrange(0, host_num):
                    src = src_sw*host_num+x
                    while (True):
                        dst_sw = random.choice(cluster_i)
                        if dst_sw != src_sw:
                            break
                    y = random.randint(0, host_num-1)
                    dst = dst_sw*host_num+y
                    #print src, src_sw, dst, dst_sw
                    matrix[src][dst] = random.random()
                    if matrix[src][dst]>threshold:
                        if src not in flows_src:
                            flows_src[src]={}
                        if dst not in flows_dst:
                            flows_dst[dst]={}
                        flows_src[src][dst] = {'converged': False, 'demand': matrix[src][dst], 'recLimited': False}
                        flows_dst[dst][src] = {'converged': False, 'demand': matrix[src][dst], 'recLimited': False}

def transform (Host_Matrix, switchList, host_num, demand_type, threshold=0.1):
    '''
    Transform the host-level traffic demand matrix to switch-level demand matrix
    '''
    base_len = len(switchList)
    sw_matrix = np.zeros((base_len, base_len), dtype=np.float16)
    starttime = datetime.now()    
    for host_src in Host_Matrix:
        switch_src = host_src/host_num
        for host_dst in Host_Matrix[host_src]:
            if Host_Matrix[host_src][host_dst]['demand'] > threshold:
                switch_dst = host_dst/host_num
                sw_matrix[switch_src-1][switch_dst-1] += Host_Matrix[host_src][host_dst]['demand']
    endtime = datetime.now()
    timediff = endtime - starttime
    print 'Transform time ', timediff
    #print sw_matrix.max()
    #print sw_matrix

    sw_matrix= sw_matrix/sw_matrix.max()
    return sw_matrix
    #print_matrix(sw_matrix, 'Traffic Demand Matrix(%s)'%demand_type, 'From Top of Rack Switch', 'To Top of Rack Switch', 'Traffic Demand Matrix(%s)'%demand_type)